package com.jxsr.algorithm;

import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Description :
 * @Author : Bruce Lee
 * @CreateTime : 2024/12/30
 */

public class TokenLimiter {

    /**
     * 令牌桶容量
     */
    private int capacity;

    /**
     * 令牌桶速率
     */
    private int rate;

    /**
     * 当前令牌数量
     */
    private AtomicInteger tokens;


    public TokenLimiter(int capacity, int rate) {
        this.capacity = capacity;
        this.rate = rate;
        this.tokens = new AtomicInteger(3);
        // 定期生成令牌
        start();
    }

    /**
     * 定期生成令牌
     */
    private void start() {

        ScheduledExecutorService executor = Executors.newScheduledThreadPool(1);

        /**
         * 每秒生成 rate 个令牌
         */
        executor.scheduleAtFixedRate(() -> {
            // 当前令牌数量
            int currentToken = tokens.get();
            if (currentToken < capacity) {
                tokens.compareAndSet(currentToken, Math.min(capacity, currentToken + rate));
            }
        }, 0, 1, java.util.concurrent.TimeUnit.SECONDS);


    }


    /**
     * 获取令牌
     */
    public boolean tryAcquire() {
        return tokens.getAndDecrement() > 0;
    }


}
